import torch
import torch.nn as tnn
from torch.utils.data import DataLoader
import torchvision.datasets as dsets
import torchvision.transforms as transforms

BATCH_SIZE = 10
LEARNING_RATE = 0.01
EPOCH = 50
N_CLASSES = 25

transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

train_set = dsets.ImageFolder('E:/Temp/vgg16/train', transform)
# test_set = dsets.ImageFolder('../data/imagenet/test', transform)

train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
# test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=True)


def conv_layer(chann_in, chann_out, k_size, p_size):
    layer = tnn.Sequential(
        tnn.Conv2d(chann_in, chann_out, kernel_size=k_size, padding=p_size),
        tnn.BatchNorm2d(chann_out),
        tnn.ReLU()
    )
    return layer


def vgg_conv_block(in_list, out_list, k_list, p_list, pooling_k, pooling_s):
    layers = [conv_layer(in_list[i], out_list[i], k_list[i], p_list[i]) for i in range(len(in_list))]
    layers += [tnn.MaxPool2d(kernel_size=pooling_k, stride=pooling_s)]
    return tnn.Sequential(*layers)


def vgg_fc_layer(size_in, size_out):
    layer = tnn.Sequential(
        tnn.Linear(size_in, size_out),
        tnn.BatchNorm1d(size_out),
        tnn.ReLU()
    )
    return layer


class VGG16(tnn.Module):
    def __init__(self, n_classes=1000):
        super(VGG16, self).__init__()

        # Conv blocks(BatchNorm + Relu activation added in each block)
        self.layer1 = vgg_conv_block([3, 64], [64, 64], [3, 3], [1, 1], 2, 2)
        self.layer2 = vgg_conv_block([64, 128], [128, 128], [3, 3], [1, 1], 2, 2)
        self.layer3 = vgg_conv_block([128, 256, 256], [256, 256, 256], [3, 3, 3], [1, 1, 1], 2, 2)
        self.layer4 = vgg_conv_block([256, 512, 512], [512, 512, 512], [3, 3, 3], [1, 1, 1], 2, 2)
        self.layer5 = vgg_conv_block([512, 512, 512], [512, 512, 512], [3, 3, 3], [1, 1, 1], 2, 2)

        self.layer6 = vgg_fc_layer(7 * 7 * 512, 4096)
        self.layer7 = vgg_fc_layer(4096, 4096)

        self.layer8 = tnn.Linear(4096, n_classes)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        vgg16_features = self.layer5(out)
        out = vgg16_features.view(out.size(0), -1)
        out = self.layer6(out)
        out = self.layer7(out)
        out = self.layer8(out)

        return vgg16_features, out


vgg16 = VGG16(n_classes=N_CLASSES)
for name, param in vgg16.named_parameters():
    print(name, param.data.shape)
vgg16.cpu()
#
# cost = tnn.CrossEntropyLoss()
# optimizer = torch.optim.Adam(vgg16.parameters(), lr=LEARNING_RATE)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)
#
# for epoch in range(EPOCH):
#     avg_loss = 0
#     cnt = 0
#     for images, labels in train_loader:
#         images = images.cpu()
#         labels = labels.cpu()
#
#         optimizer.zero_grad()
#         _, outputs = vgg16(images)
#         loss = cost(outputs, labels)
#         avg_loss += loss.data
#         cnt += 1
#         print("[E: {}] loss: {}, avg_loss: {}".format(epoch, loss.data, avg_loss / cnt))
#         loss.backward()
#         optimizer.step()
#     scheduler.step(avg_loss)

# vgg16.eval()
# correct = 0
# total = 0
#
# for images, labels in test_loader:
#     images = images.cpu()
#     _, outputs = vgg16(images)
#     _, predicted = torch.max(outputs.data, 1)
#     total += labels.size(0)
#     correct += (predicted.cpu() == labels).sum()
#     print(predicted, labels, correct, total)
#     print("avg acc: {}".format(100 * correct / total))
#
# torch.save(vgg16.state_dict(), 'cnn.pkl')
